from math import exp
import torch
import torch.nn as nn
import config as c
from rrdb_denselayer import ResidualDenseBlock_out

# 隐藏、恢复可逆块
class INV_block(nn.Module):
    def __init__(self, subnet_constructor=ResidualDenseBlock_out, clamp=c.clamp, harr=True, in_1=3, in_2=3):
        super().__init__()
        # 输入通道 * 4 
        if harr:
            self.split_len1 = in_1 * 4  # 第二行的输入通道
            self.split_len2 = in_2 * 4  # 第一行的输入通道
        # 某个常量因子
        self.clamp = clamp

        # φ模块
        self.f = subnet_constructor(self.split_len2, self.split_len1)
        # ρ模块
        self.r = subnet_constructor(self.split_len1, self.split_len2)
        # η模块
        self.y = subnet_constructor(self.split_len1, self.split_len2)

    # 图中的exp(σ(.))
    def e(self, s):
        return torch.exp(self.clamp * 2 * (torch.sigmoid(s) - 0.5))

    def forward(self, x, rev=False):
        # x.narrow函数
        # dimension：要进行切片的维度。维度从0开始计数。
        # start：切片的起始索引（包含）。
        # length：切片的长度。
        # x1表示第一行，x2表示第二行
        x1, x2 = (x.narrow(1, 0, self.split_len1),
                  x.narrow(1, self.split_len1, self.split_len2))

        # forward
        if not rev:
            # x2表示输入的秘密图像，x1表示输入的载体图像
            t2 = self.f(x2)
            # y1表示载体图像的输出
            y1 = x1 + t2
            s1, t1 = self.r(y1), self.y(y1)
            # y2表示秘密图像的输出
            y2 = self.e(s1) * x2 + t1

        # backward,恢复的过程
        else:
            # x1表示隐写图像即第二行，x2表示辅助变量z那一行
            s1, t1 = self.r(x1), self.y(x1)
            # y2 表示第一行输出
            y2 = (x2 - t1) / self.e(s1)
            t2 = self.f(y2)
            # y1表示第二行输出
            y1 = (x1 - t2)

        # 按照通道进行拼接，第二行是y1，第一行是y2，特征拼接在一起
        return torch.cat((y1, y2), 1)

